Share SQL stages and dispatch validated numerical kernels - #302

reviews:
  auto_review:
    drafts: true

• Autofix · Keep fixing CodeRabbit findings and required CI, and resolving merge conflicts

Thanks for using CodeRabbit! It's free for OSS, and your support helps us grow. If you like it, consider giving us a shout-out.

🌀 Claude: Thank you. The measurements here are excellent: 1.53× on the tiny model from computing each CTE once, with every result checked. They show clearly what recomputing CTEs costs us. This PR's concerns are about where the code lives rather than whether it works. I've opened #307 to decide where each piece should go.

Strengths

• Ordinary SQL stays the fallback, and ddx still differentiates the original views, so gradients don't depend on how the forward pass is staged.
• GPU float32 is an explicit opt-in, and the tolerances per host (host.rtol, host.atol) say so honestly.
• The tests cover NULLs, function overrides, changing inputs and snapshots.

Concerns

1. SharedExecutor is CTE materialization inside xarray-sql. That includes a SQL parser (sqlparser) in our Rust crate, a whitelist of functions, and stage caching. We decided to leave computing each CTE once to DataFusion (DataFusion plans a CTE once per reference, so ML queries grow exponentially with depth #267, on its roadmap) and to work with its developers on it, rather than carry our own.

2. xarray_sql/mlx.py is an einfold backend. It's built on einfold#156's interface and belongs with einfold, or in a small package of its own. xarray-sql would pass it through.

3. Private imports across repositories: einfold._native and ddxdb._ddxdb.to_sql, plus a subclass of ddx's DataFusionBackend. A refactor in either project would break xarray-sql.

4. The build script installs a fork. accelerators/einfold_mlx/build.py clones a personal fork at a pinned commit into the user's environment. That should become a released einfold once einfold#156 lands.

5. Demo 2's SQL. The demos are meant to be plain, legible SQL, with slowness fixed in the engine:

   • The exact_gelu, dropout and grouped_softmax flags branch the model's SQL.
   • grouped_softmax rewrites the softmax for the engine's sake.
   • The comment explaining why the max is subtracted is deleted.

   Exact GELU and dropout are fair parts of the nanoGPT recipe. They could be the demo's only spelling rather than flags.

I'd hold this PR until #307 settles. The pieces that stay in xarray-sql (the host tolerances, and the demos on public APIs) could then land on their own.